# -*- coding: utf-8 -*-
"""
车牌颜色识别模块

基于深度学习的车牌颜色分类功能，包括：
- 车牌颜色预处理
- 颜色分类模型
- 颜色识别推理

支持的车牌颜色类型：
- 黑牌、蓝牌、危险品、绿牌、白牌、黄牌

Author: AI Assistant
Date: 2025-06-12
"""

import warnings
import cv2
import torch
import numpy as np
import torch.nn as nn
from torchvision import transforms

from plate_recognition.plateNet import MyNet_color

# 车牌颜色类别定义
PLATE_COLOR_CLASSES = ['黑牌', '蓝牌', '危险品', '绿牌', '白牌', '黄牌']

# 图像标准化参数
COLOR_NORM_MEAN = [0.4243, 0.4947, 0.434]
COLOR_NORM_STD = [0.2569, 0.2478, 0.2174]

# 模型输入图像尺寸
COLOR_MODEL_INPUT_SIZE = (34, 9)

class SimplePlateColorNet(nn.Module):
    """
    简单的车牌颜色识别网络
    
    用于车牌颜色分类的轻量级卷积神经网络
    """
    
    def __init__(self, num_classes=6):
        """
        初始化网络结构
        
        Args:
            num_classes: 颜色类别数量，默认6类
        """
        super(SimplePlateColorNet, self).__init__()
        self.num_classes = num_classes
        
        self.backbone = nn.Sequential(
            # 卷积层
            nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(5, 5), stride=(1, 1)),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2)),
            nn.Dropout(0),
            
            # 全连接层
            nn.Flatten(),
            nn.Linear(480, 64),
            nn.Dropout(0),
            nn.ReLU(),
            nn.Linear(64, num_classes),
            nn.Dropout(0),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        """
        前向传播
        
        Args:
            x: 输入图像张量
            
        Returns:
            颜色分类概率分布
        """
        logits = self.backbone(x)
        return logits

def initialize_color_model(model_path, device):
    """
    初始化车牌颜色识别模型
    
    Args:
        model_path: 模型文件路径
        device: 计算设备(CPU/GPU)
        
    Returns:
        初始化后的颜色识别模型
    """
    try:
        # 忽略警告信息
        warnings.filterwarnings('ignore')
        
        # 创建模型实例
        num_classes = len(PLATE_COLOR_CLASSES)
        model = MyNet_color(num_classes)
        
        # 加载预训练权重
        state_dict = torch.load(model_path, map_location=device, weights_only=False)
        model.load_state_dict(state_dict)
        
        # 设置为评估模式并移动到指定设备
        model.eval().to(device)
        
        return model
        
    except Exception as e:
        print(f"颜色识别模型初始化失败: {e}")
        raise

def preprocess_color_image(image):
    """
    预处理车牌图像用于颜色识别
    
    Args:
        image: 输入的BGR格式车牌图像
        
    Returns:
        预处理后的图像张量
    """
    # BGR转RGB
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 调整图像尺寸
    resized_image = cv2.resize(rgb_image, COLOR_MODEL_INPUT_SIZE)
    
    # 转换为CHW格式并归一化
    image_array = np.transpose(resized_image, (2, 0, 1))
    normalized_image = image_array / 255.0
    
    # 转换为PyTorch张量
    image_tensor = torch.tensor(normalized_image, dtype=torch.float32)
    
    # 应用标准化
    normalize_transform = transforms.Normalize(mean=COLOR_NORM_MEAN, std=COLOR_NORM_STD)
    normalized_tensor = normalize_transform(image_tensor)
    
    # 添加批次维度
    batch_tensor = torch.unsqueeze(normalized_tensor, dim=0)
    
    return batch_tensor

def recognize_plate_color(plate_image, model, device):
    """
    识别车牌颜色
    
    Args:
        plate_image: 车牌图像(BGR格式)
        model: 颜色识别模型
        device: 计算设备
        
    Returns:
        识别的颜色类别名称
    """
    try:
        # 图像预处理
        input_tensor = preprocess_color_image(plate_image)
        input_tensor = input_tensor.to(device)
        
        # 模型推理
        with torch.no_grad():
            predictions = model(input_tensor)
            
        # 获取预测结果
        predicted_class_index = int(torch.argmax(predictions, dim=1)[0])
        predicted_color = PLATE_COLOR_CLASSES[predicted_class_index]
        
        return predicted_color
        
    except Exception as e:
        print(f"车牌颜色识别失败: {e}")
        return "未知颜色"

# 向后兼容的函数别名
MyNet = SimplePlateColorNet
init_color_model = initialize_color_model
plate_color_rec = recognize_plate_color

if __name__ == '__main__':
    """
    模块测试入口
    """
    try:
        # 测试配置
        test_image_path = "/mnt/Gpan/Mydata/pytorchPorject/myCrnnPlate/images/test.jpg"
        model_path = "/mnt/Gpan/Mydata/pytorchPorject/Chinese_license_plate_detection_recognition/weights/color_classify.pth"
        
        # 检查测试图像是否存在
        if not os.path.exists(test_image_path):
            print(f"测试图像不存在: {test_image_path}")
            print("请提供有效的测试图像路径")
            exit(1)
            
        if not os.path.exists(model_path):
            print(f"模型文件不存在: {model_path}")
            print("请提供有效的模型文件路径")
            exit(1)
        
        # 初始化设备
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {device}")
        
        # 加载测试图像
        test_image = cv2.imread(test_image_path)
        if test_image is None:
            print(f"无法读取测试图像: {test_image_path}")
            exit(1)
            
        print(f"测试图像尺寸: {test_image.shape}")
        
        # 初始化模型
        color_model = initialize_color_model(model_path, device)
        print("颜色识别模型初始化成功")
        
        # 执行颜色识别
        predicted_color = recognize_plate_color(test_image, color_model, device)
        
        # 输出结果
        print(f"识别结果: {predicted_color}")
        print(f"支持的颜色类别: {PLATE_COLOR_CLASSES}")
        
    except Exception as e:
        print(f"测试执行失败: {e}")
        import traceback
        traceback.print_exc()
